package com.atguigu.flink.state;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;

/**
 * Created by Smexy on 2023/3/1
 *
 *
 *  Operator State： 同一个Task的所有数据共享一个状态。
 *
 *          UnionListState: 在恢复时，是每个Task分配状态的全量
 *            和
 *          ListState:      在恢复时，使将状态 平均分配到 每个Task
 *              的区别：
 *
 *
 *          UnionListState 在KafkaSource中得到了使用。
 *              程序重启前，source有两个subTask
 *                      subTask1 ： 被分配消费 topicA-0
 *                                              [topicA-0, 20]
 *                       subTask2 ： 被分配消费 topicA-1
 *                                              [topicA-1, 21]
 *
 *            source算子会把消费者组消费的offsets信息，作为状态进行存储。
 *              程序重启，从之前的备份中恢复状态(offsets信息)，就将offsets的全量数据，分配到每一个subTask!
 *                  使用 UnionListState。
 *                      subTask1： 被分配消费 topicA-1
 *                                   [topicA-0, 20]，[topicA-1, 21]
 *                       subTask2： 被分配消费 topicA-0
 *                                   [topicA-0, 20]，[topicA-1, 21]
 *
 *
 *
 *
 */
public class Demo2_UnionListState
{
    public static void main(String[] args) {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.setParallelism(2);

        env.enableCheckpointing(2000);

                env
                   .socketTextStream("hadoop103", 8888)
                   .map(new MyMapFunction())
                   .addSink(new SinkFunction<String>()
                   {
                       //如果数据中包含了 x，手动模拟异常
                       @Override
                       public void invoke(String value, Context context) throws Exception {

                           if (value.contains("x")){
                               throw new RuntimeException("出错了...");
                           }
                           System.out.println(value);
                       }
                   });


                try {
                            env.execute();
                        } catch (Exception e) {
                            e.printStackTrace();
                        }

    }

    //累积输入的每个字符串
    public static class MyMapFunction implements MapFunction<String,String>, CheckpointedFunction
    {


        private ListState<String> listState;

        @Override
        public String map(String value) throws Exception {
            listState.add(value);
            return listState.get().toString();
        }

        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            System.out.println("MyMapFunction.snapshotState");

        }

        //初始化状态。  在Task重启时，从之前的备份中帮你恢复状态。 只在Task创建时，执行一次！
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            System.out.println("MyMapFunction.initializeState");

            listState = context.getOperatorStateStore().getUnionListState(new ListStateDescriptor<>("strs", String.class));
        }
    }
}
